"""
多因素MACD策略运行脚本
提供命令行界面，支持基础策略运行、参数优化和结果分析
"""

import pandas as pd
import numpy as np
import argparse
import os
import sys
import time
from datetime import datetime

# 导入策略类和优化器
from macd_multi_factor_strategy import MultiFactorMACDStrategy
from parameter_optimizer import ParameterOptimizer

# 设置中文字体支持
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei']
plt.rcParams['axes.unicode_minus'] = False


def create_sample_data(start_date='2023-01-01', end_date='2023-12-31', seed=42):
    """创建示例股票数据"""
    dates = pd.date_range(start_date, end_date, freq='D')
    np.random.seed(seed)
    
    # 生成带趋势的价格数据
    trend = np.linspace(100, 130, len(dates))
    noise = np.random.randn(len(dates)) * 2
    price = trend + np.cumsum(noise) * 0.4
    price = np.maximum(price, 1)  # 确保价格为正
    
    # 生成其他字段
    df = pd.DataFrame({
        '日期': dates,
        '股票代码': '000001.SZ',
        '开盘': price * (1 - np.random.uniform(0, 0.02, len(dates))),
        '收盘': price,
        '最高': price * (1 + np.random.uniform(0, 0.03, len(dates))),
        '最低': price * (1 - np.random.uniform(0, 0.03, len(dates))),
        '成交量': np.random.randint(1000000, 8000000, len(dates)),
        '成交额': price * np.random.randint(1000000, 8000000, len(dates)),
        '振幅': np.random.uniform(1, 6, len(dates)),
        '涨跌幅': np.random.uniform(-4, 4, len(dates)),
        '涨跌额': np.random.uniform(-3, 3, len(dates)),
        '换手率': np.random.uniform(0.8, 4, len(dates))
    })
    
    return df


def load_data(file_path=None):
    """加载股票数据"""
    if file_path is None:
        print("未提供数据文件路径，使用示例数据")
        return create_sample_data()
    
    try:
        # 检查文件是否存在
        if not os.path.exists(file_path):
            print(f"错误: 文件不存在: {file_path}")
            return create_sample_data()
        
        # 加载数据
        print(f"正在加载数据: {file_path}")
        df = pd.read_csv(file_path)
        
        # 检查必要字段
        required_columns = ['日期', '股票代码', '开盘', '收盘', '最高', '最低', '成交量', '成交额', '振幅', '涨跌幅', '涨跌额', '换手率']
        missing_columns = [col for col in required_columns if col not in df.columns]
        
        if missing_columns:
            print(f"警告: 缺少必要字段: {missing_columns}，可能影响策略表现")
        
        # 转换日期字段
        try:
            df['日期'] = pd.to_datetime(df['日期'])
        except Exception as e:
            print(f"警告: 日期字段转换失败: {e}")
        
        # 排序数据
        if '日期' in df.columns:
            df = df.sort_values('日期')
        df = df.reset_index(drop=True)
        
        print(f"数据加载完成: {len(df)} 行")
        
        # 数据统计
        if '日期' in df.columns and not df['日期'].empty:
            date_min = df['日期'].min().date()
            date_max = df['日期'].max().date()
            print(f"时间范围: {date_min} 至 {date_max}")
        
        if '收盘' in df.columns:
            print(f"收盘价范围: {df['收盘'].min():.2f} - {df['收盘'].max():.2f}")
        
        return df
    except Exception as e:
        print(f"数据加载失败: {e}")
        print("使用示例数据继续")
        return create_sample_data()


def run_basic_strategy(df, params, initial_capital, output_dir):
    """运行基础策略回测"""
    print("=" * 50)
    print("运行基础策略回测")
    print("=" * 50)
    
    # 创建策略实例
    strategy = MultiFactorMACDStrategy(**params)
    
    print(f"使用参数: {params}")
    print(f"初始资金: {initial_capital:,.2f}")
    
    # 记录开始时间
    start_time = time.time()
    
    # 计算指标
    print("计算技术指标...")
    df_with_indicators = df.copy()
    df_with_indicators = strategy.calculate_macd(df_with_indicators)
    df_with_indicators = strategy.calculate_rsi(df_with_indicators)
    df_with_indicators = strategy.calculate_ma_cross(df_with_indicators)
    df_with_indicators = strategy.calculate_atr(df_with_indicators)
    df_with_indicators = strategy.calculate_volume_filter(df_with_indicators)
    
    # 生成信号
    print("生成交易信号...")
    df_with_signals = strategy.generate_signals(df_with_indicators)
    
    # 执行回测
    print("执行回测...")
    results = strategy.backtest_strategy(df_with_signals, initial_capital)
    
    # 计算运行时间
    elapsed_time = time.time() - start_time
    print(f"回测完成，耗时: {elapsed_time:.2f} 秒")
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 保存结果图表
    chart_path = os.path.join(output_dir, 'strategy_results.png')
    print(f"生成结果图表: {chart_path}")
    strategy.plot_results(df_with_signals, results, save_path=chart_path)
    
    # 导出交易记录
    trade_records_path = os.path.join(output_dir, 'trade_records.csv')
    print(f"导出交易记录: {trade_records_path}")
    strategy.export_trade_records(results, trade_records_path)
    
    # 导出策略结果数据
    results_data_path = os.path.join(output_dir, 'strategy_data.csv')
    df_with_signals.to_csv(results_data_path, index=False, encoding='utf-8-sig')
    print(f"导出策略数据: {results_data_path}")
    
    print("\n" + "=" * 50)
    print("基础策略回测完成")
    print("=" * 50)
    
    return results


def run_parameter_optimization(df, complexity, initial_capital, output_dir, metric):
    """运行参数优化"""
    print("=" * 50)
    print("运行参数优化")
    print("=" * 50)
    
    # 创建优化器实例
    optimizer = ParameterOptimizer(MultiFactorMACDStrategy, initial_capital=initial_capital)
    
    # 创建参数网格
    print(f"创建{complexity}复杂度的参数网格...")
    param_grid = optimizer.create_parameter_grid(complexity=complexity)
    
    # 计算参数组合数量
    combination_count = optimizer.calculate_combination_count(param_grid)
    print(f"参数组合数量: {combination_count}")
    
    # 警告：高级复杂度可能需要很长时间
    if complexity == 'advanced' and combination_count > 1000:
        print(f"警告: 高级复杂度有 {combination_count} 个参数组合，可能需要数小时才能完成")
        confirm = input("是否继续? (y/n): ")
        if confirm.lower() != 'y':
            print("取消参数优化")
            return None
    
    # 记录开始时间
    start_time = time.time()
    
    # 运行网格搜索
    print(f"开始参数优化 (性能指标: {metric})...")
    best_params, best_performance, all_results = optimizer.run_grid_search(
        df, param_grid, performance_metric=metric, 
        use_progress_bar=True, save_results=False
    )
    
    # 计算运行时间
    elapsed_time = time.time() - start_time
    print(f"参数优化完成，耗时: {elapsed_time:.2f} 秒")
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 保存优化结果
    results_path = os.path.join(output_dir, f'optimization_results_{metric}.csv')
    optimizer.save_optimization_results(results_path)
    
    # 可视化优化结果
    viz_path = os.path.join(output_dir, f'optimization_visualization_{metric}.png')
    print(f"生成优化结果可视化: {viz_path}")
    optimizer.visualize_optimization_results(top_n=10, save_path=viz_path)
    
    # 分析关键参数敏感性
    key_params = ['fast_period', 'slow_period', 'signal_period', 'rsi_period']
    for param in key_params:
        if param in best_params:
            sens_path = os.path.join(output_dir, f'sensitivity_{param}_{metric}.png')
            print(f"生成参数敏感性分析: {sens_path}")
            try:
                optimizer.analyze_parameter_sensitivity(param, save_path=sens_path)
            except Exception as e:
                print(f"参数 {param} 敏感性分析失败: {e}")
    
    # 保存最佳参数
    best_params_path = os.path.join(output_dir, f'best_params_{metric}.txt')
    with open(best_params_path, 'w', encoding='utf-8') as f:
        f.write(f"最佳参数组合 (基于{metric}):\n")
        for k, v in best_params.items():
            f.write(f"{k}: {v}\n")
        f.write(f"\n最佳{metric}: {best_performance:.4f}\n")
    
    print(f"最佳参数已保存到: {best_params_path}")
    
    print("\n" + "=" * 50)
    print("参数优化完成")
    print("=" * 50)
    
    return best_params


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='多因素MACD策略运行脚本')
    
    # 基本参数
    parser.add_argument('--mode', type=str, default='basic', 
                        choices=['basic', 'optimize', 'best_params'],
                        help='运行模式: basic(基础回测), optimize(参数优化), best_params(使用最佳参数回测)')
    parser.add_argument('--data', type=str, default=None,
                        help='股票数据CSV文件路径，默认为示例数据')
    parser.add_argument('--capital', type=float, default=100000,
                        help='初始资金，默认为100000')
    parser.add_argument('--output', type=str, default='output',
                        help='输出目录，默认为output')
    
    # 策略参数
    parser.add_argument('--fast_period', type=int, default=12,
                        help='MACD快速均线周期，默认为12')
    parser.add_argument('--slow_period', type=int, default=26,
                        help='MACD慢速均线周期，默认为26')
    parser.add_argument('--signal_period', type=int, default=9,
                        help='MACD信号线周期，默认为9')
    parser.add_argument('--rsi_period', type=int, default=14,
                        help='RSI周期，默认为14')
    parser.add_argument('--rsi_low', type=int, default=30,
                        help='RSI超卖阈值，默认为30')
    parser.add_argument('--rsi_high', type=int, default=70,
                        help='RSI超买阈值，默认为70')
    parser.add_argument('--ma_short', type=int, default=5,
                        help='短期均线周期，默认为5')
    parser.add_argument('--ma_long', type=int, default=20,
                        help='长期均线周期，默认为20')
    parser.add_argument('--stop_loss', type=float, default=2.0,
                        help='止损ATR倍数，默认为2.0')
    parser.add_argument('--take_profit', type=float, default=3.0,
                        help='止盈ATR倍数，默认为3.0')
    
    # 参数优化参数
    parser.add_argument('--complexity', type=str, default='medium',
                        choices=['simple', 'medium', 'advanced'],
                        help='参数优化复杂度，默认为medium')
    parser.add_argument('--metric', type=str, default='sharpe_ratio',
                        choices=['sharpe_ratio', 'total_return', 'calmar_ratio'],
                        help='参数优化性能指标，默认为sharpe_ratio')
    parser.add_argument('--best_params_file', type=str, default=None,
                        help='最佳参数文件路径，用于best_params模式')
    
    return parser.parse_args()


def load_best_params(file_path):
    """加载最佳参数文件"""
    if file_path is None or not os.path.exists(file_path):
        print("未找到最佳参数文件")
        return None
    
    try:
        best_params = {}
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            for line in lines:
                line = line.strip()
                if ':' in line and not line.startswith('#') and not line.startswith('最佳'):
                    key, value = line.split(':', 1)
                    key = key.strip()
                    value = value.strip()
                    
                    # 尝试转换为适当的类型
                    try:
                        if '.' in value:
                            best_params[key] = float(value)
                        else:
                            best_params[key] = int(value)
                    except:
                        best_params[key] = value
        
        print(f"加载最佳参数: {best_params}")
        return best_params
    except Exception as e:
        print(f"最佳参数文件加载失败: {e}")
        return None


def main():
    """主函数"""
    # 解析命令行参数
    args = parse_args()
    
    # 打印运行信息
    print("多因素MACD策略运行脚本")
    print(f"运行时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"运行模式: {args.mode}")
    
    # 加载数据
    df = load_data(args.data)
    
    # 创建策略参数字典
    strategy_params = {
        'fast_period': args.fast_period,
        'slow_period': args.slow_period,
        'signal_period': args.signal_period,
        'rsi_period': args.rsi_period,
        'rsi_low_threshold': args.rsi_low,
        'rsi_high_threshold': args.rsi_high,
        'ma_short_period': args.ma_short,
        'ma_long_period': args.ma_long,
        'stop_loss_multiplier': args.stop_loss,
        'take_profit_multiplier': args.take_profit
    }
    
    # 根据运行模式执行不同的操作
    if args.mode == 'basic':
        # 运行基础策略
        run_basic_strategy(df, strategy_params, args.capital, args.output)
        
    elif args.mode == 'optimize':
        # 运行参数优化
        best_params = run_parameter_optimization(
            df, args.complexity, args.capital, args.output, args.metric
        )
        
        # 如果优化成功，可以使用最佳参数运行一次
        if best_params:
            print("\n使用最佳参数运行策略...")
            best_params_output = os.path.join(args.output, 'best_params_run')
            run_basic_strategy(df, best_params, args.capital, best_params_output)
            
    elif args.mode == 'best_params':
        # 使用最佳参数运行
        best_params = load_best_params(args.best_params_file)
        
        # 如果没有加载到最佳参数，使用命令行参数
        if best_params is None:
            print("使用命令行参数作为最佳参数")
            best_params = strategy_params
        
        # 运行策略
        run_basic_strategy(df, best_params, args.capital, args.output)
    
    print("\n" + "=" * 50)
    print("策略运行完成")
    print(f"所有结果已保存到: {args.output}")
    print("=" * 50)

if __name__ == "__main__":
    main()